#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <winsock2.h>

#pragma comment(lib, "ws2_32.lib")
#define BUFFER_SIZE 4096

int main(int argc, char *argv[]) {
    char* dir = "~";
    WSADATA wsaData;
    SOCKET sock;
    struct sockaddr_in serverAddr;
    char buffer[BUFFER_SIZE];
    int result;
    int SERVER_PORT = 4030;
	    if (argc >= 2){
	        for (int i=0;i<argc;i++) {
	            if (strcmp(argv[i], "-p") == 0) {
	                SERVER_PORT = atoi(argv[i+1]);
	            }
	        }
	    }
    WSAStartup(MAKEWORD(2, 2), &wsaData);
    sock = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);

    serverAddr.sin_family = AF_INET;
    serverAddr.sin_port = htons(SERVER_PORT);
    serverAddr.sin_addr.s_addr = inet_addr(argv[1]);

    result = connect(sock, (struct sockaddr*)&serverAddr, sizeof(serverAddr));
    if (result == SOCKET_ERROR) {
        printf("Failed to connect to server.\n");
        WSACleanup();
        return 1;
    }
    char res[BUFFER_SIZE];
    int count = 0;
    int end = 0;
    while(strcmp(res, "succN") != 0){
    	if (count >= 3) {
    		end = 1;
			break;
			
		}
		printf("%s's password:", argv[1]);
		fgets(buffer, BUFFER_SIZE, stdin);
		buffer[strcspn(buffer, "\n")] = 0;
		send(sock, buffer, strlen(buffer), 0);
		recv(sock, res, BUFFER_SIZE, 0);
		
		count++;
	}
    
    
    while (1 && end != 1) {
    	
        memset(buffer, 0, BUFFER_SIZE);
        printf("shell@%s>", argv[1]);
        fgets(buffer, BUFFER_SIZE, stdin);
        buffer[strcspn(buffer, "\n")] = 0;
        if (strcmp(buffer, "exit") == 0) {
			break;
		}
		send(sock, buffer, strlen(buffer), 0);
        memset(buffer, 0, BUFFER_SIZE);
        result = recv(sock, buffer, BUFFER_SIZE, 0);
        if (result == SOCKET_ERROR) {
            printf("Broken Pipe.\n");
            closesocket(sock);
            WSACleanup();
            return 1;
        }

        printf("%s\n", buffer);
    }
    if(end == 1) {
		printf("Login failed\n");
	}
    closesocket(sock);
    WSACleanup();
    return 0;
}
